#include "block_hhmm.hpp"

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

void metrics::report(std::ostream &os) const {
  os << "=========================================" << std::endl;
  os << "* HMM total time    : " << mat_total << " s\n";
  os << "  * HAdd time       : " << mat_add << " s\n";
  os << "  * HMM time        : " << mat_mul << " s\n";
  os << "* Final noise bound : \n";
  if (noise_m0 != 0.0) {
    os << "  @ M0 noise bound  : " << noise_m0 << std::endl;
    os << "  @ M1 noise bound  : " << noise_m1 << std::endl;
    os << "  @ M2 noise bound  : " << noise_m2 << std::endl;
    os << "  @ M3 noise bound  : " << noise_m3 << std::endl;
    os << "  @ M4 noise bound  : " << noise_m4 << std::endl;
    os << "  @ M5 noise bound  : " << noise_m5 << std::endl;
    os << "  @ M6 noise bound  : " << noise_m6 << std::endl;
  }
  for (auto i = 0; i < noise.size(); i++) {
    for (auto j = 0; j < noise[0].size(); j++) {
      os << "  * C[" << i << "][" << j << "] " << noise[i][j] << std::endl;
    }
  }
  os << "=========================================" << std::endl;
}

void blk_matmul(mat_ctxt &C, const mat_ctxt &A, const mat_ctxt &B,
                hypercube_hmme &hmme, MATRIX_TYPE type, metrics &metric) {
  metric.mat_total -= omp_get_wtime();
  metric.noise.resize(A.size());

  const helib::PubKey &pk = A[0][0].getPubKey();
  helib::Ctxt zero_ctxt(pk);
  zero_ctxt.DummyEncrypt(NTL::ZZX(0l));
  C.resize(A.size() /* m */,
           std::vector<helib::Ctxt>(B[0].size() /* n */, zero_ctxt));
  for (auto i = 0; i < A.size() /* m */; i++) {
    for (auto j = 0; j < B[0].size() /* n */; j++) {
      for (auto k = 0; k < B.size() /* k */; k++) {
        metric.mat_mul -= omp_get_wtime();
        helib::Ctxt tmp = A[i][k];
        hmme.cgemm(type, tmp, tmp, B[k][j]);
        metric.mat_mul += omp_get_wtime();
        metric.mat_add -= omp_get_wtime();
        C[i][j] += tmp;
        metric.mat_add += omp_get_wtime();
      }
      metric.noise[i].push_back(log2_noise_bound(C[i][j]));
    }
  }
  metric.mat_total += omp_get_wtime();
}

void strassen(mat_ctxt &C, const mat_ctxt &A, const mat_ctxt &B,
              hypercube_hmme &hmme, MATRIX_TYPE type, metrics &metric) {
  metric.mat_total -= omp_get_wtime();
  C.resize(A.size());
  metric.noise.resize(A.size());

  // M0 = (A_00 + A_11)(B_00 + B_11)
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M0 = A[0][0];
  helib::Ctxt tmp = B[0][0];
  M0 += A[1][1];
  tmp += B[1][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M0, M0, tmp);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m0 = log2_noise_bound(M0);

  // M1 = (A_10 + A_11)B_00
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M1 = A[1][0];
  M1 += A[1][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M1, M1, B[0][0]);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m1 = log2_noise_bound(M1);

  // M2 = A_00(B_01 - B_11)
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M2 = B[0][1];
  M2 -= B[1][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M2, A[0][0], M2);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m2 = log2_noise_bound(M2);

  // M3 = A_11(B_10 - B_00)
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M3 = B[1][0];
  M3 -= B[0][0];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M3, A[1][1], M3);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m3 = log2_noise_bound(M3);

  // M4 = (A_00 + A_01)B_11
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M4 = A[0][0];
  M4 += A[0][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M4, M4, B[1][1]);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m4 = log2_noise_bound(M4);

  // M5 = (A_10 - A_00)(B_00 + B_01)
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M5 = A[1][0];
  tmp = B[0][0];
  M5 -= A[0][0];
  tmp += B[0][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M5, M5, tmp);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m5 = log2_noise_bound(M5);

  // M6 = (A_01 - A_11)(B_10 + B_11)
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt M6 = A[0][1];
  tmp = B[1][0];
  M6 -= A[1][1];
  tmp += B[1][1];
  metric.mat_add += omp_get_wtime();
  metric.mat_mul -= omp_get_wtime();
  hmme.cgemm(type, M6, M6, tmp);
  metric.mat_mul += omp_get_wtime();
  metric.noise_m6 = log2_noise_bound(M6);

  // C_00 = M0 + M3 - M4 + M6
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt C00 = M0;
  C00 += M3;
  C00 -= M4;
  C00 += M6;
  metric.mat_add += omp_get_wtime();
  C[0].emplace_back(C00);
  metric.noise[0].push_back(log2_noise_bound(C00));

  // C_01 = M2 + M4
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt C01 = M2;
  C01 += M4;
  metric.mat_add += omp_get_wtime();
  C[0].emplace_back(C01);
  metric.noise[0].push_back(log2_noise_bound(C01));

  // C_10 = M1 + M3
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt C10 = M1;
  C10 += M3;
  metric.mat_add += omp_get_wtime();
  C[1].emplace_back(C10);
  metric.noise[1].push_back(log2_noise_bound(C10));

  // C_11 = M0 - M1 + M2 + M5
  metric.mat_add -= omp_get_wtime();
  helib::Ctxt C11 = M0;
  C11 -= M1;
  C11 += M2;
  C11 += M5;
  metric.mat_add += omp_get_wtime();
  C[1].emplace_back(C11);
  metric.noise[1].push_back(log2_noise_bound(C11));
  metric.mat_total += omp_get_wtime();
}

void strassen_ptxt(mat_ptxt &C, const mat_ptxt &A, const mat_ptxt &B) {
  C.resize(A.size());

  // M0 = (A_00 + A_11)(B_00 + B_11)
  auto M0 = (A[0][0] + A[1][1]) * (B[0][0] + B[1][1]);

  // M1 = (A_10 + A_11)B_00
  auto M1 = (A[1][0] + A[1][1]) * B[0][0];

  // M2 = A_00(B_01 - B_11)
  auto M2 = A[0][0] * (B[0][1] - B[1][1]);

  // M3 = A_11(B_10 - B_00)
  auto M3 = A[1][1] * (B[1][0] - B[0][0]);

  // M4 = (A_00 + A_01)B_11
  auto M4 = (A[0][0] + A[0][1]) * B[1][1];

  // M5 = (A_10 - A_00)(B_00 + B_01)
  auto M5 = (A[1][0] - A[0][0]) * (B[0][0] + B[0][1]);

  // M6 = (A_01 - A_11)(B_10 + B_11)
  auto M6 = (A[0][1] - A[1][1]) * (B[1][0] + B[1][1]);

  // C_00 = M0 + M3 - M4 + M6
  C[0].emplace_back(M0 + M3 - M4 + M6);

  // C_01 = M2 + M4
  C[0].emplace_back(M2 + M4);

  // C_10 = M1 + M3
  C[1].emplace_back(M1 + M3);

  // C_11 = M0 - M1 + M2 + M5
  C[1].emplace_back(M0 - M1 + M2 + M5);
}